package org.qnit.controller;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.apache.commons.lang3.StringUtils;
import org.qnit.common.MilvusTemplate;
import org.qnit.common.MysqlQuery;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.ResponseBody;

import java.util.Map;

@Controller
public class MyController {

    private Logger logger = LoggerFactory.getLogger(MyController.class);

    private String collectionName = "user_embedding";
    private String collectionName1 = "user_label_w2v";


    @Autowired
    @Qualifier("milvusTemplate")
    private MilvusTemplate milvusTemplate;

    @Autowired
    @Qualifier("mysqlQuery")
    private MysqlQuery mysqlQuery;

    @RequestMapping("/hello")
    @ResponseBody
    public String hello() {
        return "hello";
    }

    private Long toMilvusKey(String key) {
        Integer hashUserIdInt = key.hashCode();
        if (hashUserIdInt >= 0)
            return Long.valueOf("1" + hashUserIdInt);
        else
            return Long.valueOf("2" + Math.abs(hashUserIdInt));
    }

    @RequestMapping("/findUserByUser")
    @ResponseBody
    public JSONArray findUserByUser(String userId, @RequestParam(required = false, defaultValue = "10") Integer topK) {
        JSONArray respJsonArr = new JSONArray();
        try {
            Long userKey = toMilvusKey(userId);
            // milvus当余弦相识度为1时默认不返回，自身会被排除，所以topK+1
            Map<String, Float> map = milvusTemplate.searchIPVectorById(collectionName, userKey, topK + 1);
            int rank = 0;
            for (String key : map.keySet()) {
                rank++;
                JSONObject jsonObj = new JSONObject();
                Map<String, Object> mapping = mysqlQuery.getUserByMilvus(key);
                jsonObj.put("rank", rank);
                jsonObj.put("user_id", mapping.get("user_id").toString());
                jsonObj.put("score", map.get(key));
                respJsonArr.add(jsonObj);
            }
        } catch (Exception e) {
            logger.error(e.getMessage());
            respJsonArr.clear();
        } finally {
            return respJsonArr;
        }
    }

    @RequestMapping("/findUserByTerm")
    @ResponseBody
    public JSONArray findUserByTerm(String term, @RequestParam(required = false, defaultValue = "10") Integer topK) {
        JSONArray respJsonArr = new JSONArray();
        try {
            Long termKey = toMilvusKey(term);
            Map<String, Float> map = milvusTemplate.searchIPVectorById(collectionName1, collectionName, termKey, topK);
            int rank = 0;
            for (String key : map.keySet()) {
                rank++;
                JSONObject jsonObj = new JSONObject();
                Map<String, Object> mapping = mysqlQuery.getUserByMilvus(key);
                jsonObj.put("rank", rank);
                jsonObj.put("user_id", mapping.get("user_id").toString());
                jsonObj.put("score", map.get(key));
                respJsonArr.add(jsonObj);
            }
        } catch (Exception e) {
            logger.error(e.getMessage());
            respJsonArr.clear();
        } finally {
            return respJsonArr;
        }
    }

    @RequestMapping("/findUserByTermV2")
    @ResponseBody
    public JSONArray findUserByTermV2(@RequestBody JSONObject json) {
        String term = json.getString("term");
        if (StringUtils.isBlank(term))
            term = "小麦";
        Integer topK = json.getInteger("topK");
        if (topK == null)
            topK = 10;
        JSONArray respJsonArr = new JSONArray();
        try {
            Long termKey = toMilvusKey(term);
            Map<String, Float> map = milvusTemplate.searchIPVectorById(collectionName1, collectionName, termKey, topK);
            int rank = 0;
            for (String key : map.keySet()) {
                rank++;
                JSONObject jsonObj = new JSONObject();
                Map<String, Object> mapping = mysqlQuery.getUserByMilvus(key);
                jsonObj.put("label", term);
                jsonObj.put("rank", rank);
                jsonObj.put("user_id", mapping.get("user_id").toString());
                jsonObj.put("score", map.get(key));
                respJsonArr.add(jsonObj);
            }
        } catch (Exception e) {
            logger.error(e.getMessage());
            respJsonArr.clear();
        } finally {
            return respJsonArr;
        }
    }

}
